"""
敏感数据加密存储和配置信息脱敏处理
提供数据加密、解密、配置脱敏等安全功能
"""
import os
import json
import base64
import hashlib
import secrets
from typing import Any, Optional, Dict, Union, List
from datetime import datetime, timedelta
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend
from loguru import logger

from utils.logger import Logger
from utils.config import Config
from utils.exceptions import SecurityException, ErrorCode


class EncryptionManager:
    """加密管理器"""
    
    def __init__(self):
        self.config = Config()
        self.logger = Logger.get_logger(agent_name="encryption_manager")
        
        # 加密配置
        self.encryption_config = {
            "algorithm": self.config.get("security.encryption.algorithm", "fernet"),
            "key_rotation_enabled": self.config.get("security.encryption.key_rotation_enabled", True),
            "key_rotation_interval": self.config.get("security.encryption.key_rotation_interval", 86400),  # 24小时
            "master_key_path": self.config.get("security.encryption.master_key_path", "keys/master.key"),
            "backup_keys_path": self.config.get("security.encryption.backup_keys_path", "keys/backup/")
        }
        
        # 主密钥
        self.master_key: Optional[bytes] = None
        self.current_key: Optional[bytes] = None
        self.key_creation_time: Optional[datetime] = None
        
        # 初始化加密管理器
        self._initialize_encryption()
        
        self.logger.info("加密管理器初始化完成")
    
    def _initialize_encryption(self):
        """初始化加密系统"""
        try:
            # 加载或生成主密钥
            self.master_key = self._load_or_generate_master_key()
            
            # 生成当前加密密钥
            self.current_key = self._generate_encryption_key()
            self.key_creation_time = datetime.now()
            
            # 启动密钥轮换任务
            if self.encryption_config["key_rotation_enabled"]:
                self._start_key_rotation_task()
            
        except Exception as e:
            self.logger.error(f"加密系统初始化失败: {str(e)}")
            raise SecurityException(
                error_code=ErrorCode.ENCRYPTION_ERROR,
                message="加密系统初始化失败"
            )
    
    def _load_or_generate_master_key(self) -> bytes:
        """加载或生成主密钥"""
        key_path = self.encryption_config["master_key_path"]
        
        # 检查密钥文件是否存在
        if os.path.exists(key_path):
            try:
                with open(key_path, 'rb') as f:
                    return f.read()
            except Exception as e:
                self.logger.warning(f"主密钥文件读取失败，生成新密钥: {str(e)}")
        
        # 生成新主密钥
        master_key = Fernet.generate_key()
        
        # 确保目录存在
        os.makedirs(os.path.dirname(key_path), exist_ok=True)
        
        # 保存主密钥
        try:
            with open(key_path, 'wb') as f:
                f.write(master_key)
            
            # 设置文件权限（仅所有者可读写）
            os.chmod(key_path, 0o600)
            
            self.logger.info("主密钥生成并保存成功")
            return master_key
            
        except Exception as e:
            self.logger.error(f"主密钥保存失败: {str(e)}")
            raise SecurityException(
                error_code=ErrorCode.ENCRYPTION_ERROR,
                message="主密钥保存失败"
            )
    
    def _generate_encryption_key(self) -> bytes:
        """生成加密密钥"""
        # 使用主密钥派生加密密钥
        salt = secrets.token_bytes(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        
        key = base64.urlsafe_b64encode(kdf.derive(self.master_key))
        return key
    
    def _start_key_rotation_task(self):
        """启动密钥轮换任务"""
        # 这里应该实现一个后台任务来定期轮换密钥
        # 为了简化，这里只是记录日志
        self.logger.info("密钥轮换任务已启动")
    
    def encrypt_data(self, data: Union[str, bytes, dict]) -> str:
        """
        加密数据
        
        Args:
            data: 要加密的数据
            
        Returns:
            加密后的数据（base64编码）
        """
        try:
            # 将数据转换为字节
            if isinstance(data, dict):
                data_bytes = json.dumps(data).encode('utf-8')
            elif isinstance(data, str):
                data_bytes = data.encode('utf-8')
            elif isinstance(data, bytes):
                data_bytes = data
            else:
                data_bytes = str(data).encode('utf-8')
            
            # 创建Fernet实例
            f = Fernet(self.current_key)
            
            # 加密数据
            encrypted_bytes = f.encrypt(data_bytes)
            
            # 返回base64编码的字符串
            encrypted_str = base64.b64encode(encrypted_bytes).decode('utf-8')
            
            self.logger.debug("数据加密成功")
            return encrypted_str
            
        except Exception as e:
            self.logger.error(f"数据加密失败: {str(e)}")
            raise SecurityException(
                error_code=ErrorCode.ENCRYPTION_ERROR,
                message="数据加密失败"
            )
    
    def decrypt_data(self, encrypted_data: str) -> Union[str, dict]:
        """
        解密数据
        
        Args:
            encrypted_data: 加密的数据（base64编码）
            
        Returns:
            解密后的数据
        """
        try:
            # 解码base64
            encrypted_bytes = base64.b64decode(encrypted_data.encode('utf-8'))
            
            # 创建Fernet实例
            f = Fernet(self.current_key)
            
            # 解密数据
            decrypted_bytes = f.decrypt(encrypted_bytes)
            
            # 尝试解析为JSON
            try:
                decrypted_str = decrypted_bytes.decode('utf-8')
                return json.loads(decrypted_str)
            except json.JSONDecodeError:
                # 如果不是JSON，返回字符串
                return decrypted_bytes.decode('utf-8')
            
        except Exception as e:
            self.logger.error(f"数据解密失败: {str(e)}")
            raise SecurityException(
                error_code=ErrorCode.DECRYPTION_ERROR,
                message="数据解密失败"
            )
    
    def hash_data(self, data: Union[str, bytes], algorithm: str = "sha256") -> str:
        """
        哈希数据
        
        Args:
            data: 要哈希的数据
            algorithm: 哈希算法
            
        Returns:
            哈希值（十六进制字符串）
        """
        try:
            # 转换数据为字节
            if isinstance(data, str):
                data_bytes = data.encode('utf-8')
            elif isinstance(data, bytes):
                data_bytes = data
            else:
                data_bytes = str(data).encode('utf-8')
            
            # 选择哈希算法
            if algorithm.lower() == "md5":
                hash_obj = hashlib.md5()
            elif algorithm.lower() == "sha1":
                hash_obj = hashlib.sha1()
            elif algorithm.lower() == "sha256":
                hash_obj = hashlib.sha256()
            elif algorithm.lower() == "sha512":
                hash_obj = hashlib.sha512()
            else:
                raise ValueError(f"不支持的哈希算法: {algorithm}")
            
            # 计算哈希
            hash_obj.update(data_bytes)
            return hash_obj.hexdigest()
            
        except Exception as e:
            self.logger.error(f"数据哈希失败: {str(e)}")
            raise SecurityException(
                error_code=ErrorCode.INTERNAL_ERROR,
                message="数据哈希失败"
            )
    
    def generate_secure_token(self, length: int = 32) -> str:
        """
        生成安全令牌
        
        Args:
            length: 令牌长度
            
        Returns:
            安全令牌
        """
        return secrets.token_urlsafe(length)


class DataMasker:
    """数据脱敏器"""
    
    def __init__(self):
        self.config = Config()
        self.logger = Logger.get_logger(agent_name="data_masker")
        
        # 脱敏配置
        self.masking_config = {
            "enabled": self.config.get("security.masking.enabled", True),
            "log_masking": self.config.get("security.masking.log_masking", True),
            "api_masking": self.config.get("security.masking.api_masking", True),
            "mask_char": self.config.get("security.masking.mask_char", "*"),
            "email_show_prefix": self.config.get("security.masking.email_show_prefix", 3),
            "phone_show_prefix": self.config.get("security.masking.phone_show_prefix", 3),
            "phone_show_suffix": self.config.get("security.masking.phone_show_suffix", 4),
            "id_card_show_prefix": self.config.get("security.masking.id_card_show_prefix", 4),
            "id_card_show_suffix": self.config.get("security.masking.id_card_show_suffix", 4),
            "credit_card_show_prefix": self.config.get("security.masking.credit_card_show_prefix", 6),
            "credit_card_show_suffix": self.config.get("security.masking.credit_card_show_suffix", 4)
        }
        
        # 敏感字段配置
        self.sensitive_fields = {
            "password": {"type": "full", "description": "密码"},
            "token": {"type": "partial", "show_chars": 4, "description": "访问令牌"},
            "secret": {"type": "full", "description": "密钥"},
            "private_key": {"type": "full", "description": "私钥"},
            "api_key": {"type": "partial", "show_chars": 8, "description": "API密钥"},
            "email": {"type": "email", "description": "邮箱地址"},
            "phone": {"type": "phone", "description": "手机号码"},
            "mobile": {"type": "phone", "description": "手机号码"},
            "id_card": {"type": "id_card", "description": "身份证号码"},
            "idcard": {"type": "id_card", "description": "身份证号码"},
            "credit_card": {"type": "credit_card", "description": "信用卡号"},
            "bank_card": {"type": "credit_card", "description": "银行卡号"},
            "bank_account": {"type": "credit_card", "description": "银行账户"}
        }
        
        self.logger.info("数据脱敏器初始化完成")
    
    def mask_data(self, data: Any, field_name: Optional[str] = None, 
                  mask_type: str = "auto") -> Any:
        """
        脱敏数据
        
        Args:
            data: 要脱敏的数据
            field_name: 字段名（用于自动检测脱敏类型）
            mask_type: 脱敏类型
            
        Returns:
            脱敏后的数据
        """
        if not self.masking_config["enabled"]:
            return data
        
        try:
            if isinstance(data, dict):
                return self._mask_dict(data)
            elif isinstance(data, list):
                return [self.mask_data(item, field_name, mask_type) for item in data]
            elif isinstance(data, str):
                return self._mask_string(data, field_name, mask_type)
            else:
                return data
                
        except Exception as e:
            self.logger.error(f"数据脱敏失败: {str(e)}")
            return "****"  # 出错时返回默认脱敏值
    
    def _mask_dict(self, data: Dict[str, Any]) -> Dict[str, Any]:
        """脱敏字典数据"""
        masked_data = {}
        
        for key, value in data.items():
            # 检测字段类型
            field_type = self._detect_field_type(key)
            
            if field_type:
                masked_data[key] = self._mask_value(value, field_type)
            else:
                # 递归处理嵌套数据
                masked_data[key] = self.mask_data(value, key)
        
        return masked_data
    
    def _detect_field_type(self, field_name: str) -> Optional[str]:
        """检测字段类型"""
        field_name_lower = field_name.lower()
        
        # 检查字段名是否包含敏感关键词
        for sensitive_field, config in self.sensitive_fields.items():
            if sensitive_field in field_name_lower:
                return sensitive_field
        
        return None
    
    def _mask_value(self, value: Any, field_type: str) -> str:
        """脱敏值"""
        if not isinstance(value, str):
            value = str(value)
        
        config = self.sensitive_fields.get(field_type, {"type": "full"})
        mask_type = config["type"]
        
        if mask_type == "full":
            return self.masking_config["mask_char"] * min(len(value), 8)
        
        elif mask_type == "partial":
            show_chars = config.get("show_chars", 4)
            if len(value) > show_chars:
                return value[:show_chars] + self.masking_config["mask_char"] * 4
            return self.masking_config["mask_char"] * min(len(value), 8)
        
        elif mask_type == "email":
            return self._mask_email(value)
        
        elif mask_type == "phone":
            return self._mask_phone(value)
        
        elif mask_type == "id_card":
            return self._mask_id_card(value)
        
        elif mask_type == "credit_card":
            return self._mask_credit_card(value)
        
        else:
            return self.masking_config["mask_char"] * min(len(value), 8)
    
    def _mask_string(self, value: str, field_name: Optional[str] = None, 
                     mask_type: str = "auto") -> str:
        """脱敏字符串"""
        if mask_type == "auto" and field_name:
            field_type = self._detect_field_type(field_name)
            if field_type:
                return self._mask_value(value, field_type)
        
        return value
    
    def _mask_email(self, email: str) -> str:
        """脱敏邮箱"""
        if "@" not in email:
            return self.masking_config["mask_char"] * min(len(email), 8)
        
        username, domain = email.split("@", 1)
        show_prefix = self.masking_config["email_show_prefix"]
        
        if len(username) > show_prefix:
            masked_username = username[:show_prefix] + self.masking_config["mask_char"] * 3
        else:
            masked_username = self.masking_config["mask_char"] * min(len(username), 3)
        
        return f"{masked_username}@{domain}"
    
    def _mask_phone(self, phone: str) -> str:
        """脱敏手机号"""
        # 移除非数字字符
        digits = ''.join(filter(str.isdigit, phone))
        
        if len(digits) < 7:
            return self.masking_config["mask_char"] * min(len(phone), 8)
        
        show_prefix = self.masking_config["phone_show_prefix"]
        show_suffix = self.masking_config["phone_show_suffix"]
        
        masked_digits = (digits[:show_prefix] + 
                        self.masking_config["mask_char"] * (len(digits) - show_prefix - show_suffix) +
                        digits[-show_suffix:])
        
        return masked_digits
    
    def _mask_id_card(self, id_card: str) -> str:
        """脱敏身份证"""
        # 移除非数字字符
        digits = ''.join(filter(str.isdigit, id_card))
        
        if len(digits) < 15:
            return self.masking_config["mask_char"] * min(len(id_card), 8)
        
        show_prefix = self.masking_config["id_card_show_prefix"]
        show_suffix = self.masking_config["id_card_show_suffix"]
        
        masked_digits = (digits[:show_prefix] + 
                        self.masking_config["mask_char"] * (len(digits) - show_prefix - show_suffix) +
                        digits[-show_suffix:])
        
        return masked_digits
    
    def _mask_credit_card(self, credit_card: str) -> str:
        """脱敏信用卡"""
        # 移除非数字字符
        digits = ''.join(filter(str.isdigit, credit_card))
        
        if len(digits) < 8:
            return self.masking_config["mask_char"] * min(len(credit_card), 8)
        
        show_prefix = self.masking_config["credit_card_show_prefix"]
        show_suffix = self.masking_config["credit_card_show_suffix"]
        
        masked_digits = (digits[:show_prefix] + 
                        self.masking_config["mask_char"] * (len(digits) - show_prefix - show_suffix) +
                        digits[-show_suffix:])
        
        return masked_digits


class SecureConfigManager:
    """安全配置管理器"""
    
    def __init__(self):
        self.config = Config()
        self.logger = Logger.get_logger(agent_name="secure_config_manager")
        self.encryption_manager = EncryptionManager()
        self.data_masker = DataMasker()
        
        # 敏感配置项
        self.sensitive_config_keys = [
            "database.password",
            "redis.password",
            "api.secret_key",
            "api.jwt_secret",
            "email.password",
            "sms.api_key",
            "storage.secret_key",
            "payment.private_key",
            "oauth.client_secret"
        ]
        
        self.logger.info("安全配置管理器初始化完成")
    
    def get_secure_config(self, key: str, default: Any = None, 
                         mask_for_logging: bool = True) -> Any:
        """
        获取安全配置
        
        Args:
            key: 配置键
            default: 默认值
            mask_for_logging: 是否对日志进行脱敏
            
        Returns:
            配置值
        """
        try:
            # 检查是否是敏感配置
            is_sensitive = key in self.sensitive_config_keys
            
            # 获取原始配置值
            value = self.config.get(key, default)
            
            # 如果是敏感配置，尝试解密
            if is_sensitive and isinstance(value, str):
                try:
                    # 检查是否是加密值（通过前缀判断）
                    if value.startswith("ENC(") and value.endswith(")"):
                        encrypted_value = value[4:-1]  # 移除ENC()前缀和后缀
                        value = self.encryption_manager.decrypt_data(encrypted_value)
                except Exception as e:
                    self.logger.warning(f"配置解密失败，使用原始值: {key}, 错误: {str(e)}")
            
            # 记录日志（脱敏）
            if mask_for_logging and is_sensitive:
                masked_value = self.data_masker.mask_data(str(value), key)
                self.logger.debug(f"获取安全配置: {key} = {masked_value}")
            else:
                self.logger.debug(f"获取配置: {key} = {value}")
            
            return value
            
        except Exception as e:
            self.logger.error(f"获取配置失败: {key}, 错误: {str(e)}")
            return default
    
    def set_secure_config(self, key: str, value: Any, 
                         encrypt_if_sensitive: bool = True) -> bool:
        """
        设置安全配置
        
        Args:
            key: 配置键
            value: 配置值
            encrypt_if_sensitive: 如果是敏感配置是否加密
            
        Returns:
            是否成功
        """
        try:
            is_sensitive = key in self.sensitive_config_keys
            
            # 加密敏感配置
            if is_sensitive and encrypt_if_sensitive and value is not None:
                try:
                    encrypted_value = self.encryption_manager.encrypt_data(str(value))
                    value = f"ENC({encrypted_value})"
                    self.logger.debug(f"配置已加密: {key}")
                except Exception as e:
                    self.logger.error(f"配置加密失败: {key}, 错误: {str(e)}")
                    return False
            
            # 设置配置（这里应该调用配置管理器的设置方法）
            # 由于Config类可能没有set方法，这里只是记录日志
            self.logger.info(f"配置设置: {key} = {'***' if is_sensitive else value}")
            return True
            
        except Exception as e:
            self.logger.error(f"配置设置失败: {key}, 错误: {str(e)}")
            return False
    
    def mask_config_for_display(self, config_dict: Dict[str, Any]) -> Dict[str, Any]:
        """
        为显示脱敏配置
        
        Args:
            config_dict: 配置字典
            
        Returns:
            脱敏后的配置字典
        """
        masked_config = {}
        
        for key, value in config_dict.items():
            # 检查是否是敏感配置
            is_sensitive = any(sensitive_key in key for sensitive_key in [
                "password", "secret", "key", "token", "private"
            ])
            
            if is_sensitive and value is not None:
                masked_config[key] = self.data_masker.mask_data(str(value), key)
            else:
                masked_config[key] = value
        
        return masked_config
    
    def validate_config_security(self, config_dict: Dict[str, Any]) -> Dict[str, bool]:
        """
        验证配置安全性
        
        Args:
            config_dict: 配置字典
            
        Returns:
            安全验证结果
        """
        security_checks = {
            "has_weak_passwords": False,
            "has_plain_text_secrets": False,
            "has_expired_keys": False,
            "has_insecure_algorithms": False
        }
        
        try:
            # 检查弱密码
            password_fields = ["database.password", "redis.password", "email.password"]
            for field in password_fields:
                password = config_dict.get(field, "")
                if password and len(password) < 8:
                    security_checks["has_weak_passwords"] = True
                    self.logger.warning(f"检测到弱密码: {field}")
            
            # 检查明文密钥
            secret_fields = ["api.secret_key", "api.jwt_secret", "payment.private_key"]
            for field in secret_fields:
                secret = config_dict.get(field, "")
                if secret and not secret.startswith("ENC("):
                    security_checks["has_plain_text_secrets"] = True
                    self.logger.warning(f"检测到明文密钥: {field}")
            
            # 检查过期密钥（这里只是示例，实际应该有密钥有效期检查）
            # ...
            
            # 检查不安全的算法
            algorithm_fields = ["security.encryption.algorithm", "api.signature_algorithm"]
            for field in algorithm_fields:
                algorithm = config_dict.get(field, "")
                if algorithm and algorithm.lower() in ["md5", "sha1", "des"]:
                    security_checks["has_insecure_algorithms"] = True
                    self.logger.warning(f"检测到不安全算法: {field} = {algorithm}")
            
            return security_checks
            
        except Exception as e:
            self.logger.error(f"配置安全验证失败: {str(e)}")
            return security_checks


# 全局实例
_encryption_manager: Optional[EncryptionManager] = None
_data_masker: Optional[DataMasker] = None
_secure_config_manager: Optional[SecureConfigManager] = None


def get_encryption_manager() -> EncryptionManager:
    """获取加密管理器实例"""
    global _encryption_manager
    if _encryption_manager is None:
        _encryption_manager = EncryptionManager()
    return _encryption_manager


def get_data_masker() -> DataMasker:
    """获取数据脱敏器实例"""
    global _data_masker
    if _data_masker is None:
        _data_masker = DataMasker()
    return _data_masker


def get_secure_config_manager() -> SecureConfigManager:
    """获取安全配置管理器实例"""
    global _secure_config_manager
    if _secure_config_manager is None:
        _secure_config_manager = SecureConfigManager()
    return _secure_config_manager